import torch
import os
import numpy as np
from PIL import Image
from typing import List, Dict, Any, Union, Tuple
from cn_clip.clip import load_from_name, available_models
import torchvision.transforms as transforms

class ClipModel:
    """CN_CLIP模型封装类，用于提取图片特征"""
    
    def __init__(self, model_name: str = "ViT-B-16", device: str = None):
        """
        初始化CN_CLIP模型
        
        Args:
            model_name: 模型名称，默认为'ViT-B-16'
            device: 设备，默认为自动检测
        """
        if device is None:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device
        
        print(f"加载CN_CLIP模型 {model_name}...")
        self.model_name = model_name
        self.model, self.preprocess = load_from_name(model_name, device=self.device)
        print(f"模型已加载到{self.device}")
        
    @staticmethod
    def get_available_models() -> List[str]:
        """获取可用的CN_CLIP模型列表"""
        return available_models()
    
    def extract_features(self, image_paths: List[str]) -> Tuple[np.ndarray, List[str]]:
        """
        从图片中提取特征向量
        
        Args:
            image_paths: 图片路径列表
            
        Returns:
            特征向量和对应的图片路径
        """
        valid_images = []
        valid_paths = []
        
        for img_path in image_paths:
            try:
                image = Image.open(img_path).convert('RGB')
                valid_images.append(self.preprocess(image))
                valid_paths.append(img_path)
            except Exception as e:
                print(f"警告: 无法处理图片 {img_path}: {str(e)}")
        
        if not valid_images:
            return np.array([]), []
        
        # 将图片转换为批处理张量
        image_batch = torch.stack(valid_images).to(self.device)
        
        # 提取特征
        with torch.no_grad():
            image_features = self.model.encode_image(image_batch)
            image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        
        # 将特征转为NumPy数组
        features = image_features.cpu().numpy()
        return features, valid_paths
    
    def extract_text_features(self, texts: List[str]) -> np.ndarray:
        """
        从文本中提取特征向量
        
        Args:
            texts: 文本列表
            
        Returns:
            文本特征向量
        """
        if not texts:
            return np.array([])
            
        # 将文本转为张量
        try:
            with torch.no_grad():
                # 检查模型是否有tokenize方法
                if hasattr(self.model, 'tokenize'):
                    text_tokens = torch.cat([self.model.tokenize(text) for text in texts]).to(self.device)
                    text_features = self.model.encode_text(text_tokens)
                    text_features = text_features / text_features.norm(dim=-1, keepdim=True)
                    
                    # 将特征转为NumPy数组
                    features = text_features.cpu().numpy()
                    return features
                else:
                    # 如果没有tokenize方法，使用备用方法
                    print("警告: 模型不支持文本编码，返回空特征")
                    return np.array([])
        except Exception as e:
            print(f"文本特征提取失败: {str(e)}")
            return np.array([])
    
    def encode_text(self, text: str) -> np.ndarray:
        """
        将单个文本转换为特征向量
        
        Args:
            text: 文本字符串
            
        Returns:
            文本特征向量
        """
        if not text:
            return None
        
        # 将文本转为张量
        try:
            with torch.no_grad():
                # 检查模型是否有tokenize方法
                if hasattr(self.model, 'tokenize'):
                    text_tokens = self.model.tokenize(text).to(self.device)
                    text_features = self.model.encode_text(text_tokens)
                    text_features = text_features / text_features.norm(dim=-1, keepdim=True)
                    
                    # 将特征转为NumPy数组
                    features = text_features.cpu().numpy()
                    return features
                else:
                    # 如果没有tokenize方法，使用备用方法
                    return self.extract_text_features([text])
        except Exception as e:
            print(f"文本编码失败: {str(e)}")
            # 返回空数组而不是None，保持一致性
            return np.array([]) 